import java.util.Arrays;

class Solution {
  public int[] getLeastNumbers(int[] arr, int k) {
    if (k == 0) {
      return new int[0];
    } else if (arr.length <= k) {
      return arr;
    }

    partArray(arr, 0, arr.length - 1, k);

    int[] res = new int[k];
    for (int i = 0; i < k; i++) {
      res[i] = arr[i];
    }
    return res;
  }

  private void partArray(int[] arr, int left, int right, int k) {
    int pivot = partition(arr, left, right);
    if (k == pivot) {
      return;
    } else if (pivot > k) {
      partArray(arr, left, pivot - 1, k);
    } else {
      partArray(arr, pivot + 1, right, k);
    }
  }

  private int partition(int[] arr, int left, int right) {
    int mid = (left + right) >> 1;
    if (arr[mid] > arr[right]) {
      swap(arr, mid, right);
    }
    if (arr[left] > arr[right]) {
      swap(arr, left, right);
    }
    if (arr[mid] > arr[left]) {
      swap(arr, mid, left);
    }
    int pivot = arr[left];
    while (left < right) {
      while (left < right && arr[right] >= pivot) {
        right--;
      }
      arr[left] = arr[right];
      while (left < right && arr[left] <= pivot) {
        left++;
      }
      arr[right] = arr[left];
    }
    arr[left] = pivot;
    return left;
  }

  private void swap(int[] arr, int i, int j) {
    int temp = arr[i];
    arr[i] = arr[j];
    arr[j] = temp;
  }

  public static void main(String[] args) {
    int[] arr = new Solution().getLeastNumbers(new int[] { 3, 2, 1 }, 2);
    System.out.println(Arrays.toString(arr));
  }
}